import os
import pandas as pd
from sqlalchemy import create_engine, text
from sqlalchemy.exc import SQLAlchemyError

# 数据库配置
DB_CONFIG = {
    'host': 'localhost',
    'database': 'books_db',
    'user': 'root',
    'password': '******',
    'port': 3306,
    'charset': 'utf8mb4'
}

engine = create_engine(
    f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}?charset={DB_CONFIG['charset']}",
    pool_pre_ping=True
)

TABLE_NAME = 'books'
CSV_FOLDER = '../cleaned'

# 创建表SQL
CREATE_TABLE_SQL = f"""
CREATE TABLE IF NOT EXISTS {TABLE_NAME} (
    id INT PRIMARY KEY AUTO_INCREMENT,
    title VARCHAR(255) NOT NULL,
    price DECIMAL(10, 2) NOT NULL,
    author VARCHAR(255) NOT NULL,
    publisher VARCHAR(255) NOT NULL,
    comments INT NOT NULL
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
"""

# 默认值设置
DEFAULT_PUBLISHER = '未知出版社'
DEFAULT_PRICE = 0.0
DEFAULT_AUTHOR = '佚名'
DEFAULT_TITLE = '未知书名'  # 新增标题默认值


def create_table():
    """创建数据表"""
    try:
        with engine.connect() as connection:
            connection.execute(text(CREATE_TABLE_SQL))
            connection.commit()
        print(f"表 {TABLE_NAME} 创建成功")
    except SQLAlchemyError as e:
        print(f"创建表失败: {str(e)}")
        raise


def import_csv_files():
    """导入多个CSV文件到数据库，处理所有缺失值"""
    csv_files = [f for f in os.listdir(CSV_FOLDER) if f.endswith('.csv')]

    if not csv_files:
        print("未找到CSV文件")
        return

    total_records = 0
    filled_title = 0
    filled_author = 0
    filled_publisher = 0
    filled_price = 0

    for csv_file in csv_files:
        file_path = os.path.join(CSV_FOLDER, csv_file)
        try:
            df = pd.read_csv(file_path)

            column_mapping = {
                '书名': 'title',
                '价格': 'price',
                '作者': 'author',
                '出版社': 'publisher',
                '评论数': 'comments'
            }

            # 确保所有列存在
            for col in column_mapping.keys():
                if col not in df.columns:
                    df[col] = ''

            df = df[list(column_mapping.keys())].rename(columns=column_mapping)

            # 处理标题列（关键修改点）
            original_title_null = df['title'].isnull().sum()
            empty_title_mask = (df['title'].str.strip() == '')

            # 填充空值和空字符串
            df['title'] = df['title'].fillna(DEFAULT_TITLE)
            df.loc[empty_title_mask, 'title'] = DEFAULT_TITLE

            filled_title += original_title_null + empty_title_mask.sum()

            # 处理作者列
            original_author_null = df['author'].isnull().sum()
            empty_author_mask = (df['author'].str.strip() == '')
            df['author'] = df['author'].fillna(DEFAULT_AUTHOR)
            df.loc[empty_author_mask, 'author'] = DEFAULT_AUTHOR
            filled_author += original_author_null + empty_author_mask.sum()

            # 处理价格列
            original_price_null = df['price'].isnull().sum()
            df['price'] = pd.to_numeric(df['price'], errors='coerce')
            new_price_null = df['price'].isnull().sum() - original_price_null
            df['price'] = df['price'].fillna(DEFAULT_PRICE)
            filled_price += original_price_null + new_price_null

            # 处理评论数列
            df['comments'] = pd.to_numeric(df['comments'], errors='coerce').fillna(0).astype(int)

            # 处理出版社
            original_pub_null = df['publisher'].isnull().sum()
            df['publisher'] = df['publisher'].fillna(DEFAULT_PUBLISHER)
            empty_pub_mask = (df['publisher'].str.strip() == '')
            df.loc[empty_pub_mask, 'publisher'] = DEFAULT_PUBLISHER
            filled_publisher += original_pub_null + empty_pub_mask.sum()

            # 导入数据
            df.to_sql(
                name=TABLE_NAME,
                con=engine,
                if_exists='append',
                index=False,
                chunksize=1000
            )

            records_imported = len(df)
            total_records += records_imported
            print(
                f"成功导入 {csv_file}: {records_imported} 条记录，填充了 {filled_title} 个标题、{filled_author} 个作者、{filled_publisher} 个出版社和 {filled_price} 个价格")

        except Exception as e:
            print(f"处理文件 {csv_file} 时出错: {str(e)}")
            continue

    print(
        f"全部导入完成，共导入 {total_records} 条记录，填充了 {filled_title} 个标题、{filled_author} 个作者、{filled_publisher} 个出版社和 {filled_price} 个价格")


if __name__ == "__main__":
    try:
        create_table()
        import_csv_files()
    except Exception as e:
        print(f"执行过程中发生错误: {str(e)}")
    finally:
        if 'engine' in locals():
            engine.dispose()
            print("数据库连接已关闭")